#include "spinlock.h"
#include "assert.h"
#include "interrupt.h"
#include "atom.h"

/**
 * initialize the spin lock
 * @param pspinlock: the descriptor of spin lock
 * @return: void
 */
void mx_spinlock_init (mx_spinlock_pt pspinlock)
{
    pspinlock->flag = MX_ATOM_FLAG_UNLOCK;
}

/**
 * lock the specified spin lock, keep trying until success
 * @param pspinlock: the descriptor of spin lock
 * @return: void
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
void mx_spinlock_lock (mx_spinlock_pt pspinlock)
{
    assert(pspinlock == mx_null);

    while (mx_true)
    {
        mx_usize_t istate = mx_interrupt_disable_nest();
        /* try lock */
        if (mx_atom_set8test((mx_usize_pt)&pspinlock->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* get lock successful, and record the local intr state for unlock */
            pspinlock->istate = istate;
            break;
        }
        /* other core locks failed, but the interrupt still has the chance
            to be executed */
        mx_interrupt_enable_nest(istate);
    }
}

/**
 * try lock the specified spin lock
 * @param pspinlock: the descriptor of spin lock
 * @return: the try lock operation is successful or not
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
mx_bool_t mx_spinlock_trylock (mx_spinlock_pt pspinlock)
{
    assert(pspinlock == mx_null);

    mx_bool_t ret = mx_false;
    mx_usize_t istate = mx_interrupt_disable_nest();
    /* try lock */
    if (mx_atom_set8test((mx_usize_pt)&pspinlock->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
    {
        pspinlock->istate = istate;
        ret = mx_true;
    }
    else
    {
        mx_interrupt_enable_nest(istate);
    }
    return ret;
}

/**
 * unlock the specified spin lock
 * @param pspinlock: the descriptor of spin lock
 * @return: void
 * @note: only the lock owner can do unlock operation
 */
void mx_spinlock_unlock (mx_spinlock_pt pspinlock)
{
    assert(pspinlock == mx_null);

    /* force unlocking */
    pspinlock->flag = MX_ATOM_FLAG_UNLOCK;
    /* enable the local interrupt */
    mx_interrupt_enable_nest(pspinlock->istate);
}

/// ------------------------------------------------------------------------ rw spinlock

/* state of rw spinlock */
#define MX_SPINLOCK_RW_WRITING MX_USIZE_MAX
#define MX_SPINLOCK_RW_UNLOCKED 0U

/**
 * rw-spinlock initialize
 * @param psrw: the descriptor of rw-spinlock
 * @return: void
 */
void mx_spinlock_rw_init (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    psrw->rw = 0;
    psrw->flag = MX_ATOM_FLAG_UNLOCK;
}

/**
 * get the read lock of rw-spinlock
 * @param psrw: the descriptor of rw-spinlock
 * @return: void
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
void mx_spinlock_rw_rlock (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    while (mx_true)
    {
        mx_usize_t istate = mx_interrupt_disable_nest();
        /* try to access `rw` */
        if (mx_atom_set8test((mx_usize_pt)&psrw->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* now we can access `rw` and check the state(`rw`) of rw-spinlock */
            if (psrw->rw != MX_SPINLOCK_RW_WRITING)
            {
                /* the current state of rwlock is no-locks or only read-locks,
                   then cumulative read lock's count value */
                psrw->rw ++;
                assert(psrw->rw == MX_SPINLOCK_RW_WRITING); // to many readers to overlap `MX_SPINLOCK_RW_WRITING`
                /* save the istate for unlock */
                psrw->istates[mx_cpu_getid()] = istate;
                /* after access `rw` is completed, exit the flag-lock so that
                   others can access it */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                break;
            }
            /* fail to get rlock then allow other access `rw` to release lock,
               need do it before enable interrupt */
            psrw->flag = MX_ATOM_FLAG_UNLOCK;
        }
        /* other core locks failed, but the interrupt still has the chance
           to be executed */
        mx_interrupt_enable_nest(istate);
    }
}

/**
 * try to get the read lock of rw-spinlock
 * @param psrw: the descriptor of rw-spinlock
 * @return: whether get the read lock successful or not
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
mx_bool_t mx_spinlock_rw_tryrlock (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    mx_bool_t ret = mx_false;
    while (mx_true)
    {
        mx_usize_t istate = mx_interrupt_disable_nest();
        /* try to access `rw` */
        if (mx_atom_set8test((mx_usize_pt)&psrw->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* now we can access `rw` and check the state(`rw`) of rw-spinlock */
            if (psrw->rw != MX_SPINLOCK_RW_WRITING)
            {
                /* the current state of rwlock is no-locks or only read-locks,
                   then cumulative read lock's count value */
                psrw->rw ++;
                assert(psrw->rw == MX_SPINLOCK_RW_WRITING); // to many readers to overlap `MX_SPINLOCK_RW_WRITING`
                /* save the istate for unlock */
                psrw->istates[mx_cpu_getid()] = istate;
                /* access `rw` over and unlock flag */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                ret = mx_true;
            }
            else
            {
                /* access `rw` over and unlock flag, need do it before enable interrupt */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                mx_interrupt_enable_nest(istate);
            }
            break;
        }
        mx_interrupt_enable_nest(istate);
    }
    return ret;
}

/**
 * get the write lock of rw-spinlock
 * @param psrw: the descriptor of rw-spinlock
 * @return: void
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
void mx_spinlock_rw_wlock (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    while (mx_true)
    {
        mx_usize_t istate = mx_interrupt_disable_nest();
        /* try to access `rw` */
        if (mx_atom_set8test((mx_usize_pt)&psrw->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* now we can access `rw` and check the state(`rw`) of rw-spinlock */
            if (psrw->rw == MX_SPINLOCK_RW_UNLOCKED)
            {
                /* get read-lock need no reader and writer, and then mark locked */
                psrw->rw = MX_SPINLOCK_RW_WRITING;
                /* save the istate for unlock */
                psrw->istates[mx_cpu_getid()] = istate;
                /* after access `rw` is completed, exit the flag-lock so that
                   others can access it */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                break;
            }
            /* fail to get rlock then allow other access `rw` to release lock */
            psrw->flag = MX_ATOM_FLAG_UNLOCK;
        }
        /* other core locks failed, but the interrupt still has the chance
           to be executed */
        mx_interrupt_enable_nest(istate);
    }
}

/**
 * try to get the write lock of rw-spinlock
 * @param psrw: the descriptor of rw-spinlock
 * @return: whether get the write lock successful or not
 * @note: if the operation is successful, it will disable the local interrupt
 * @note: self-nesting of lock is forbidden
 */
mx_bool_t mx_spinlock_rw_trywlock (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    mx_bool_t ret = mx_false;
    while (mx_true)
    {
        mx_usize_t istate = mx_interrupt_disable_nest();
        /* try to access `rw` */
        if (mx_atom_set8test((mx_usize_pt)&psrw->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* now we can access `rw` and check the state(`rw`) of rw-spinlock */
            if (psrw->rw == MX_SPINLOCK_RW_UNLOCKED)
            {
                /* get read-lock need no reader and writer, and then mark locked */
                psrw->rw = MX_SPINLOCK_RW_WRITING;
                /* save the istate for unlock */
                psrw->istates[mx_cpu_getid()] = istate;
                /* access `rw` over and unlock flag */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                ret = mx_true;
            }
            else
            {
                 /* access `rw` over and unlock flag, need do it before enable interrupt */
                psrw->flag = MX_ATOM_FLAG_UNLOCK;
                mx_interrupt_enable_nest(istate);
            }
            break;
        }
        mx_interrupt_enable_nest(istate);
    }
    return ret;
}

/**
 * unlock the rw-spinlock, both rlock and wlock can call this funciton to unlock
 * @param psrw: the descriptor of rw-spinlock
 * @return: void
 * @note: only the lock owner can do unlock operation
 */
void mx_spinlock_rw_unlock (mx_spinlock_rw_pt psrw)
{
    assert(psrw == mx_null);

    while (mx_true)
    {
        /* only the locke owner can do unlock operation, so the interrupt must be disabled.
           try to access `rw` */
        if (mx_atom_set8test((mx_usize_pt)&psrw->flag, MX_ATOM_FLAG_LOCKED, MX_ATOM_FLAG_UNLOCK))
        {
            /* check the type of the lock (rlock or wlock) */
            if (psrw->rw == MX_SPINLOCK_RW_WRITING)
            {
                /* unlock the write lock */
                psrw->rw = MX_SPINLOCK_RW_UNLOCKED;
            }
            else
            {
                /* if used in pairs (xlock and unlock), the `rw` cann't be `MX_SPINLOCK_RW_UNLOCKED(0)` */
                assert(psrw->rw == MX_SPINLOCK_RW_UNLOCKED);
                /* unlock the read lock, reader -1 */
                psrw->rw --;
            }
            /* access `rw` over and unlock flag, need do it before enable interrupt */
            psrw->flag = MX_ATOM_FLAG_UNLOCK;
            /* restore the interrupt */
            mx_interrupt_enable_nest(psrw->istates[mx_cpu_getid()]);
            break;
        }
    }
}
